import functools
from pprint import pprint

import pandas as pd
import numpy as np

def reduce_test():
    df = pd.DataFrame({'AAA': [4, 5, 6, 7], 'BBB': [10, 20, 30, 40], 'CCC': [100, 50, -30, -50]})
    print(df)
    Crit1 = df.AAA <= 5.5
    Crit2 = df.BBB == 10.0
    Crit3 = df.CCC > -40.0
    AllCrit = Crit1 & Crit2 & Crit3
    # 打印筛选结果
    print(df[AllCrit])

    # 另外一种方法 - reduce
    CritList = [Crit1, Crit2, Crit3]
    AllCrit2 = functools.reduce(lambda x, y: x & y, CritList)
    print(AllCrit2)
    print(df[AllCrit2])

def selection_test():
    # create with DatetimeIndex
    dates = pd.date_range('20130101', periods=8)
    # print(dates)
    data = np.arange(0, 24, 1).reshape(8, 3)
    # print(data)
    df = pd.DataFrame(data, index=dates, columns=list('ABC'))
    print(df)
    # get A
    print("获取一列:")
    print(df.A)  # or df['A']
    # get row
    row = df.loc[dates[0]]
    print("获取一行")
    print(row)
    print("获取单独的数据")
    print(row[['A']], type(row[['A']]))
    print(row['A'], type(row['A']))

def new_cols_test():
    df = pd.DataFrame({'AAA': [1, 2, 1, 3], 'BBB': [1, 1, 2, 2], 'CCC': [2, 1, 3, 1]})
    print("df:\n", df)

    source_cols = df.columns
    new_cols = [str(x) + "_cat" for x in source_cols]
    categories = {1: 'Alpha', 2: 'Beta', 3: 'Charlie'}
    df[new_cols] = df[source_cols].applymap(categories.get)
    print("new df:\n", df)

def idxmin_test():
    df = pd.DataFrame( {'AAA': [1, 1, 1, 2, 2, 2, 3, 3], 'BBB': [2, 1, 3, 4, 5, 1, 2, 3]})
    print("df:\n", df)
    index_min = df.groupby('AAA')["BBB"].idxmin()
    print("idxmin:\n", index_min, type(index_min), index_min.index)
    print("df.loc:\n", df.loc[index_min])
    print("df.sort_values:\n", df.sort_values(by="BBB").groupby('AAA', as_index=False).first())

if __name__ == '__main__':
    # reduce_test()
    # selection_test()
    # new_cols_test()
    idxmin_test()